from __future__ import annotations

import logging
import os
from typing import List, Optional, Union, Dict, Any, Tuple
import numpy as np
from scipy.spatial.distance import cosine

import requests
from dotenv import load_dotenv
load_dotenv()

logger = logging.getLogger(__name__)

# ------------------ 自定义异常 ------------------
class EmbeddingError(RuntimeError):
    """任何与 embedding 调用相关的异常均继承自此基类。"""

class HTTPError(EmbeddingError):
    """HTTP 层面的错误（非 2xx）。"""

class TokenLimitError(EmbeddingError):
    """输入超出模型 token 限制。"""

# ------------------ 核心封装 ------------------
class EmbeddingClient:
    """
    一个极简但健壮的 embedding 客户端，支持多种模型和格式。

    用法示例
    --------
    >>> client = EmbeddingClient(
    ...     api_key="YOUR_API_KEY",
    ...     base_url="https://api.siliconflow.cn/v1/embeddings",
    ...     model="BAAI/bge-large-zh-v1.5"
    ... )
    >>> vec = client.embed_single("Hello world")
    >>> print(len(vec))
    1024
    
    # 使用不同的编码格式
    >>> client = EmbeddingClient(
    ...     api_key="YOUR_API_KEY",
    ...     model="BAAI/bge-large-zh-v1.5",
    ...     encoding_format="base64"
    ... )
    >>> encoded_vec = client.embed_single("Hello world")
    >>> print(type(encoded_vec))
    <class 'str'>
    
    # 使用Qwen模型并指定维度
    >>> client = EmbeddingClient(
    ...     api_key="YOUR_API_KEY",
    ...     model="Qwen/Qwen3-Embedding-4B",
    ...     dimensions=512
    ... )
    >>> vec = client.embed_single("Hello world")
    >>> print(len(vec))
    512
    
    # 计算两个文本之间的相似度
    >>> client = EmbeddingClient(api_key="YOUR_API_KEY")
    >>> similarity = client.compute_similarity("苹果手机", "iPhone")
    >>> print(similarity)
    0.85
    
    # 查找与查询文本最相似的内容
    >>> client = EmbeddingClient(api_key="YOUR_API_KEY")
    >>> texts = ["苹果手机", "苹果水果", "香蕉", "电脑"]
    >>> results = client.find_similar("iPhone", texts, top_k=2)
    >>> for result in results:
    ...     print(f"{result['text']}: {result['similarity']:.2f}")
    苹果手机: 0.85
    电脑: 0.65
    """

    def __init__(
        self,
        api_key: Optional[str] = None,
        *,
        base_url: str = "https://api.siliconflow.cn/v1/embeddings",
        model: str = "BAAI/bge-large-zh-v1.5",
        timeout: float = 30.0,
        encoding_format: str = "float",
        dimensions: Optional[int] = None,
    ):
        """
        Parameters
        ----------
        api_key : str, optional
            Silicon Flow 访问令牌。也可通过环境变量 `SILICON_API_KEY` 传入。
        base_url : str
            embedding 接口完整 URL。
        model : str
            要使用的模型名称。
        timeout : float
            单次请求超时（秒）。
        encoding_format : str, optional
            返回嵌入向量的格式，可选值为"float"或"base64"，默认为"float"。
        dimensions : int, optional
            输出嵌入向量的维度，仅在使用Qwen/Qwen3系列模型时支持。
            - Qwen/Qwen3-Embedding-8B: [64,128,256,512,768,1024,2048,4096]
            - Qwen/Qwen3-Embedding-4B: [64,128,256,512,768,1024,2048]
            - Qwen/Qwen3-Embedding-0.6B: [64,128,256,512,768,1024]
        """
        self._api_key = api_key or os.getenv("SILICONCLOUD_API_KEY")
        if not self._api_key:
            raise ValueError("api_key 未提供且未找到环境变量 SILICONCLOUD_API_KEY")

        self._base_url = base_url
        self._model = model
        self._timeout = timeout
        self._encoding_format = encoding_format
        self._dimensions = dimensions
        self._session = requests.Session()
        self._session.headers.update(
            {
                "Authorization": f"Bearer {self._api_key}",
                "Content-Type": "application/json",
            }
        )

    def embed_single(self, text: str) -> Union[List[float], str]:
        """
        对一段文本生成 embedding 向量。

        Returns
        -------
        List[float] or str
            当encoding_format为"float"时，返回浮点数列表，长度取决于模型和dimensions参数。
            当encoding_format为"base64"时，返回base64编码的字符串。

        Raises
        ------
        EmbeddingError
            任何上游或下游错误均会被转换为 EmbeddingError 的子类。
        """
        return self.embed_batch([text])[0]

    def embed_batch(self, texts: List[str]) -> Union[List[List[float]], List[str]]:
        """
        批量生成 embedding，节省网络往返。

        Parameters
        ----------
        texts : List[str]
            输入文本列表，长度不可超过模型单次最大 token 限制。
            不同模型的最大token限制如下：
            - BAAI/bge-large-zh-v1.5, BAAI/bge-large-en-v1.5, netease-youdao/bce-embedding-base_v1: 512
            - BAAI/bge-m3, Pro/BAAI/bge-m3: 8192
            - Qwen/Qwen3-Embedding-8B, Qwen/Qwen3-Embedding-4B, Qwen/Qwen3-Embedding-0.6B: 32768

        Returns
        -------
        List[List[float]] or List[str]
            与输入顺序一一对应的向量列表。
            当encoding_format为"float"时，返回浮点数列表的列表。
            当encoding_format为"base64"时，返回base64编码的字符串列表。
        """
        payload = {"model": self._model, "input": texts}
        
        # 添加可选参数
        if self._encoding_format != "float":
            payload["encoding_format"] = self._encoding_format
            
        if self._dimensions is not None:
            payload["dimensions"] = self._dimensions
        try:
            resp = self._session.post(
                self._base_url, json=payload, timeout=self._timeout
            )
            resp.raise_for_status()
        except requests.HTTPError as e:
            raise HTTPError(f"HTTP {e.response.status_code}: {e.response.text}") from e
        except requests.RequestException as e:
            raise EmbeddingError("Network error") from e

        data = resp.json()
        # 校验格式
        if "data" not in data or not isinstance(data["data"], list):
            raise EmbeddingError("Unexpected response format")

        vectors = [item["embedding"] for item in data["data"]]
        logger.debug("embed_batch success, %d vectors returned", len(vectors))
        return vectors
        
    def find_similar(
        self, 
        query: str, 
        texts: List[str], 
        pre_computed_embeddings: Optional[List[List[float]]] = None,
        top_k: Optional[int] = None
    ) -> List[Dict[str, Any]]:
        """
        查找与查询文本最相似的内容。

        Parameters
        ----------
        query : str
            查询文本。
        texts : List[str]
            文本集合，用于与查询文本比较相似度。
        pre_computed_embeddings : Optional[List[List[float]]], optional
            预先计算好的文本集合的嵌入向量。如果提供，将直接使用而不重新计算。
            必须与texts列表一一对应。
        top_k : Optional[int], optional
            返回的最相似结果数量。如果为None，则返回所有结果。

        Returns
        -------
        List[Dict[str, Any]]
            按相似度降序排列的结果列表，每个结果包含以下字段：
            - text: 原始文本
            - similarity: 相似度分数（0-1之间，越大表示越相似）
            - index: 在原始texts列表中的索引
        """
        # 获取查询文本的嵌入向量
        query_embedding = self.embed_single(query)
        
        # 如果没有提供预计算的嵌入向量，则计算文本集合的嵌入向量
        if pre_computed_embeddings is None:
            corpus_embeddings = self.embed_batch(texts)
        else:
            if len(pre_computed_embeddings) != len(texts):
                raise ValueError("预计算的嵌入向量数量必须与文本数量相同")
            corpus_embeddings = pre_computed_embeddings
        
        # 计算相似度
        results = []
        for i, embedding in enumerate(corpus_embeddings):
            # 计算余弦相似度（1 - 余弦距离）
            if isinstance(query_embedding, str) or isinstance(embedding, str):
                # 如果是base64编码，需要先解码
                if self._encoding_format == "base64":
                    raise ValueError("base64编码格式不支持直接计算相似度，请使用float格式")
            
            similarity = 1 - cosine(query_embedding, embedding)
            results.append({
                "text": texts[i],
                "similarity": float(similarity),  # 转换为Python原生float，确保可JSON序列化
                "index": i
            })
        
        # 按相似度降序排序
        results.sort(key=lambda x: x["similarity"], reverse=True)
        
        # 如果指定了top_k，则只返回前top_k个结果
        if top_k is not None:
            results = results[:top_k]
            
        return results
        
    def compute_similarity(self, text1: str, text2: str) -> float:
        """
        计算两个文本之间的相似度。

        Parameters
        ----------
        text1 : str
            第一个文本。
        text2 : str
            第二个文本。

        Returns
        -------
        float
            相似度分数（0-1之间，越大表示越相似）。

        Raises
        ------
        ValueError
            如果使用base64编码格式，则无法直接计算相似度。
        """
        if self._encoding_format == "base64":
            raise ValueError("base64编码格式不支持直接计算相似度，请使用float格式")
            
        # 获取两个文本的嵌入向量
        embedding1 = self.embed_single(text1)
        embedding2 = self.embed_single(text2)
        
        # 计算余弦相似度（1 - 余弦距离）
        similarity = 1 - cosine(embedding1, embedding2)
        
        return float(similarity)  # 转换为Python原生float，确保可JSON序列化
        
    def compute_similarities_batch(
        self, 
        texts1: List[str], 
        texts2: Optional[List[str]] = None
    ) -> Union[List[List[float]], List[float]]:
        """
        批量计算文本之间的相似度。

        Parameters
        ----------
        texts1 : List[str]
            第一组文本列表。
        texts2 : Optional[List[str]], optional
            第二组文本列表。如果为None，则计算texts1中所有文本两两之间的相似度。

        Returns
        -------
        Union[List[List[float]], List[float]]
            如果提供了texts2，则返回一个相似度矩阵，其中matrix[i][j]表示texts1[i]和texts2[j]之间的相似度。
            如果没有提供texts2，则返回一个相似度列表，其中list[i]表示texts1[i]和texts1[i+1:]中每个文本之间的相似度
            （即上三角矩阵的展平形式，不包括对角线）。

        Raises
        ------
        ValueError
            如果使用base64编码格式，则无法直接计算相似度。
        """
        if self._encoding_format == "base64":
            raise ValueError("base64编码格式不支持直接计算相似度，请使用float格式")
            
        # 获取文本的嵌入向量
        embeddings1 = self.embed_batch(texts1)
        
        if texts2 is not None:
            # 计算两组文本之间的相似度矩阵
            embeddings2 = self.embed_batch(texts2)
            similarities = []
            for emb1 in embeddings1:
                row_similarities = []
                for emb2 in embeddings2:
                    similarity = 1 - cosine(emb1, emb2)
                    row_similarities.append(float(similarity))
                similarities.append(row_similarities)
            return similarities
        else:
            # 计算同一组文本内部的相似度（上三角矩阵）
            similarities = []
            for i in range(len(embeddings1)):
                for j in range(i + 1, len(embeddings1)):
                    similarity = 1 - cosine(embeddings1[i], embeddings1[j])
                    similarities.append(float(similarity))
            return similarities
            
    def sort_by_similarity(
        self, 
        query: str, 
        texts: List[str], 
        pre_computed_embeddings: Optional[List[List[float]]] = None,
        include_scores: bool = True
    ) -> Union[List[str], List[Tuple[str, float]]]:
        """
        根据与查询文本的相似度对文本列表进行排序。

        Parameters
        ----------
        query : str
            查询文本。
        texts : List[str]
            要排序的文本列表。
        pre_computed_embeddings : Optional[List[List[float]]], optional
            预先计算好的文本集合的嵌入向量。如果提供，将直接使用而不重新计算。
            必须与texts列表一一对应。
        include_scores : bool, optional
            是否在结果中包含相似度分数，默认为True。

        Returns
        -------
        Union[List[str], List[Tuple[str, float]]]
            按相似度降序排列的文本列表。
            如果include_scores为True，则返回(文本, 相似度分数)元组的列表。
            如果include_scores为False，则仅返回文本列表。

        Raises
        ------
        ValueError
            如果使用base64编码格式，则无法直接计算相似度。
        """
        # 使用find_similar方法获取排序后的结果
        results = self.find_similar(query, texts, pre_computed_embeddings)
        
        # 根据include_scores参数决定返回格式
        if include_scores:
            return [(result["text"], result["similarity"]) for result in results]
        else:
            return [result["text"] for result in results]
            
    def sort_by_precomputed_similarity(
        self, 
        similarities: Union[List[List[float]], List[float]],
        texts1: List[str],
        texts2: Optional[List[str]] = None,
        include_scores: bool = True
    ) -> Union[List[Tuple[str, str, float]], List[Tuple[str, str]]]:
        """
        根据预先计算的相似度对文本对进行排序。

        Parameters
        ----------
        similarities : Union[List[List[float]], List[float]]
            预先计算的相似度，可以是矩阵或列表形式。
            - 如果是矩阵形式，则应该是compute_similarities_batch(texts1, texts2)的结果。
            - 如果是列表形式，则应该是compute_similarities_batch(texts1)的结果。
        texts1 : List[str]
            第一组文本列表。
        texts2 : Optional[List[str]], optional
            第二组文本列表。如果为None，则假设similarities是texts1内部的相似度。
        include_scores : bool, optional
            是否在结果中包含相似度分数，默认为True。

        Returns
        -------
        Union[List[Tuple[str, str, float]], List[Tuple[str, str]]]
            按相似度降序排列的文本对列表。
            如果include_scores为True，则返回(文本1, 文本2, 相似度分数)元组的列表。
            如果include_scores为False，则返回(文本1, 文本2)元组的列表。
        """
        if texts2 is not None:
            # 处理两组文本之间的相似度矩阵
            if not isinstance(similarities, list) or not all(isinstance(row, list) for row in similarities):
                raise ValueError("当提供texts2时，similarities应该是一个矩阵")
                
            if len(similarities) != len(texts1) or any(len(row) != len(texts2) for row in similarities):
                raise ValueError("similarities矩阵的维度与texts1和texts2的长度不匹配")
                
            # 将矩阵展平为三元组列表：(text1, text2, similarity)
            pairs = []
            for i, row in enumerate(similarities):
                for j, sim in enumerate(row):
                    pairs.append((texts1[i], texts2[j], sim))
            
            # 按相似度降序排序
            pairs.sort(key=lambda x: x[2], reverse=True)
            
            # 根据include_scores参数决定返回格式
            if include_scores:
                return pairs
            else:
                return [(pair[0], pair[1]) for pair in pairs]
        else:
            # 处理同一组文本内部的相似度列表（上三角矩阵）
            if not isinstance(similarities, list) or isinstance(similarities[0], list):
                raise ValueError("当不提供texts2时，similarities应该是一个一维列表")
                
            # 计算预期的相似度列表长度
            expected_length = len(texts1) * (len(texts1) - 1) // 2
            if len(similarities) != expected_length:
                raise ValueError(f"similarities列表的长度与texts1的长度不匹配，预期长度为{expected_length}")
                
            # 重建文本对和相似度的对应关系
            pairs = []
            idx = 0
            for i in range(len(texts1)):
                for j in range(i + 1, len(texts1)):
                    pairs.append((texts1[i], texts1[j], similarities[idx]))
                    idx += 1
            
            # 按相似度降序排序
            pairs.sort(key=lambda x: x[2], reverse=True)
            
            # 根据include_scores参数决定返回格式
            if include_scores:
                return pairs
            else:
                return [(pair[0], pair[1]) for pair in pairs]